package cn.sccdlg.lucene;

import org.apache.lucene.analysis.Analyzer;
import org.apache.lucene.analysis.TokenStream;
import org.apache.lucene.document.Document;
import org.apache.lucene.index.DirectoryReader;
import org.apache.lucene.queryparser.classic.QueryParser;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.Query;
import org.apache.lucene.search.TopDocs;
import org.apache.lucene.search.highlight.*;
import org.apache.lucene.store.Directory;
import org.apache.lucene.store.FSDirectory;
import org.junit.Test;
import org.wltea.analyzer.lucene.IKAnalyzer;

import java.io.File;

/**
 * @author RaoShuang
 * @emil 991207823@qq.com
 * @date 2020/10/31 15:32
 * 索引库查询结果高亮显示
 */
public class HighlighterSearch {

    @Test
    public void testTightSearch() throws Exception {

        // 设置要查询的索引目录
        Directory directory = FSDirectory.open(new File("D:/test/index").toPath());
        // 创建indexSearcher
        DirectoryReader dReader = DirectoryReader.open(directory);
        IndexSearcher searcher = new IndexSearcher(dReader);
        // 设置分词方式
        Analyzer analyzer = new IKAnalyzer();
        // 设置查询域
        String field = "content";
        QueryParser parser = new QueryParser(field, analyzer);
        // 查询字符串
        Query query = parser.parse("lucene是一个Java开发的全文检索工具包");
        QueryScorer scorer = new QueryScorer(query, field);
        SimpleHTMLFormatter fors = new SimpleHTMLFormatter("<span style=\"color:red;\">", "</span>");
        Highlighter highlighter = new Highlighter(fors, scorer);

        System.out.println("query:" + query.toString());
        // 返回前10条
        TopDocs topDocs = searcher.search(query, 10);
        if (topDocs != null) {
            System.out.println("符合条件第文档总数：" + topDocs.totalHits);

            for (int i = 0; i < topDocs.scoreDocs.length; i++) {
                Document doc = searcher.doc(topDocs.scoreDocs[i].doc);
                TokenStream tokenStream = TokenSources.getTokenStream(searcher.getIndexReader(), topDocs.scoreDocs[i].doc, field, analyzer);
                Fragmenter fragment = new SimpleSpanFragmenter(scorer);
                highlighter.setTextFragmenter(fragment);
                //高亮news_title域
                String str = highlighter.getBestFragment(tokenStream, doc.get("content"));//获取高亮的片段，可以对其数量进行限制
                System.out.println("高亮title：" + str);
                tokenStream = TokenSources.getTokenStream(searcher.getIndexReader(), topDocs.scoreDocs[i].doc, "content", analyzer);
                str = highlighter.getBestFragment(tokenStream, doc.get("content"));//获取高亮的片段，可以对其数量进行限制
                System.out.println("高亮summary：" + str);
            }
        }
        directory.close();
        dReader.close();
    }
}
